// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/buffer/buffer.h>

namespace kumo::pollux::dwio::common {
    /// Concatenates multiple bit vectors and exposes the result in a
    /// single Buffer. Creates and resizes the result Buffer as needed. If
    /// only one bits were added, sets the output buffer to nullptr.
    class BitConcatenation {
    public:
        explicit BitConcatenation(memory::MemoryPool &pool) : pool_(pool) {
        }

        /// Prepares to concatenate bits given to append() or appendOnes()
        /// into 'buffer'. 'buffer' is allocated and resized as needed. If
        /// 'buffer' is initially nullptr and only ones are appended, buffer
        /// may stay nullptr. The size() of 'buffer' is set to the next byte,
        /// so the caller must  use numBits() to get the bit count.
        void reset(BufferPtr &outBuffer) {
            buffer_ = &outBuffer;
            numBits_ = 0;
            hasZeros_ = false;
        }

        /// Appends  'bits' between bit offset 'begin' and 'end' to the result.
        /// A nullptr 'bits' is treated as a bit range with all bits set.
        void append(const uint64_t *bits, int32_t begin, int32_t end);

        /// Appends 'numOnes' ones.
        void appendOnes(int32_t numOnes);

        // Returns 'buffer_' or nullptr if only ones have been appended.
        BufferPtr buffer() const {
            return hasZeros_ ? *buffer_ : nullptr;
        }

        int32_t numBits() const {
            return numBits_;
        }

    private:
        // Allocates or reallocates '*buffer' to have space for 'numBits_ + newBits'
        // bits. Retuns a pointer to the first word of 'buffer_'.
        uint64_t *ensureSpace(int32_t newBits);

        void setSize() {
            if (*buffer_) {
                (*buffer_)->setSize(bits::divRoundUp(numBits_, 8));
            }
        }

        memory::MemoryPool &pool_;
        BufferPtr *buffer_{nullptr};
        int32_t numBits_{0};
        bool hasZeros_{false};
    };
} // namespace kumo::pollux::dwio::common
